import os
import cv2
import numpy as np
from myCVHelper import my_cv_helper as helper
from myCVHelper.my_cv_helper import logger
from PIL import Image, ImageEnhance
import traceback
import matplotlib
from matplotlib import pyplot as plt
import random

# 设置matplotlib正常显示中文和负号
matplotlib.rcParams['font.sans-serif'] = ['SimHei']  # 用黑体显示中文
matplotlib.rcParams['axes.unicode_minus'] = False  # 正常显示负号


class Configs:
    file_path = './imgs'
    file_exts = ['jpg', ]
    window_name = "process_"


def process_pre_tasks(im: np.ndarray, show=False) -> np.ndarray:
    helper.Show.window_clear()
    helper.Show.resize = 4
    helper.Show.offset = (200, 100)
    if show:
        helper.Show.imshow(Configs.window_name + "原图", im, 0)
    # 转换到hsv处理
    im = cv2.cvtColor(im, cv2.COLOR_BGR2HLS)
    if show:
        helper.Show.imshow(Configs.window_name + "HSV", im, 0)
    # 中值滤波降噪
    im = cv2.medianBlur(im, 15)
    if show:
        helper.Show.imshow(Configs.window_name + "中值滤波", im, 0)
    # 消除高光
    _, mask = cv2.threshold(im, 100, 255, cv2.THRESH_BINARY)
    if show:
        helper.Show.imshow(Configs.window_name + "mask", mask, 0)
    im = cv2.illuminationChange(im, mask)
    if show:
        helper.Show.imshow(Configs.window_name + "消除高光", im, 1)

    # 增加对比度
    image = Image.fromarray(im)
    image = ImageEnhance.Contrast(image).enhance(3)
    im = np.array(image)
    if show:
        helper.Show.imshow(Configs.window_name + "增强对比度", im, 1)

    if show:
        helper.Controls.wait_exit(1)
    return im


def process(src: np.ndarray) -> np.ndarray:
    im = process_pre_tasks(src)
    helper.Show.window_clear()
    helper.Show.resize = 2
    helper.Show.offset = (20, 10)
    # 转灰度
    im = cv2.cvtColor(im, cv2.COLOR_HLS2BGR)
    im = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)
    helper.Show.imshow(Configs.window_name + "灰度", im, 0)

    # 二值化
    _, im = cv2.threshold(im, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    helper.Show.imshow(Configs.window_name + "二值化", im, 0)

    im2 = im.copy()

    # 开操作
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (31, 31))
    im = cv2.morphologyEx(im, cv2.MORPH_OPEN, kernel)
    helper.Show.imshow(Configs.window_name + "开操作", im, 0)

    # 描边
    im = cv2.Canny(im, 50, 200)
    helper.Show.imshow(Configs.window_name + "边缘检测", im, 1)

    # 直线检测
    lines = cv2.HoughLinesP(im, 1, np.pi / 180, 60, minLineLength=50, maxLineGap=30)
    # print('lines:', lines)
    # 绘图
    im = cv2.cvtColor(im, cv2.COLOR_GRAY2BGR)
    if lines is None: lines = []
    for line in lines:
        cv2.line(im, (line[0][0], line[0][1]), (line[0][2], line[0][3]),
                 (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)), 15)
    helper.Show.imshow(Configs.window_name + "直线", im, 1)

    # 计算k值，画出直方图
    def k_val(line_):
        return (line_[0][1] - line_[0][3]) / (line_[0][0] - line_[0][2]) if line_[0][0] - line_[0][2] != 0 else np.inf

    ks = list(map(k_val, lines))
    # 做个截断
    # ks = [k for k in ks if abs(k) != np.inf]
    ks = [k if abs(k) != np.inf else 100 for k in ks]
    ks = [k if k != 0 else 0.01 for k in ks]
    # 聚类一下
    theta_d = 20
    theta = list(map(lambda x: np.arctan(x) * 360 / np.pi, ks))
    # 每d度分一组，然后删掉小于等于两条线的组
    theta_group = [[] for _ in range(-190, 190, theta_d)]
    p = 0
    for i in range(-190, 190, theta_d):
        for t in theta:
            if i < t <= i + theta_d:
                theta_group[p].append(t)
        if len(lines > 10) and len(theta_group[p]) <= 1:
            theta_group[p] = []
        p += 1
    theta_group_avg = [sum(tg) / len(tg) if len(tg) != 0 else None for tg in theta_group]
    theta_div = [tg for tg in theta_group_avg if tg is not None]
    # 再合并一次
    theta_div2 = []
    theta_i_ignores = []
    for i in range(len(theta_div)):
        if i in theta_i_ignores:
            continue
        if i == len(theta_div) - 1:
            theta_div2.append(theta_div[i])
            break
        if theta_div[i + 1] - theta_div[i] <= theta_d \
                or abs(theta_div[i + 1] - theta_div[i]) + theta_d >= 360:
            theta_div2.append((theta_div[i] + theta_div[i + 1]) / 2)
            theta_i_ignores.append(i + 1)
        else:
            theta_div2.append(theta_div[i])

    # print('theta_div:', theta_div)
    logger.debug('theta_div2: %s' % str(theta_div2))

    reach_cross = False

    if len(theta_div2) >= 2:
        reach_cross = True
        logger.info('Reach cross!')

    # plt.hist(theta, bins=40, facecolor="red", edgecolor="green", alpha=0.7)
    # # 显示横轴标签
    # plt.xlabel("区间")
    # plt.xlim((-180, 180))
    # # 显示纵轴标签
    # plt.ylabel("频数/频率")
    # # 显示图标题
    # plt.title("频数/频率分布直方图")
    # plt.show()

    # 画出边缘
    # # 对于每个角度，选择斜率相近的直线，然后选取这些直线的端点做平均，再过这个点做直线
    # ks = list(map(np.tan, theta_div2))
    # ks_select = [[] for _ in ks]
    # k_d = 0.1
    # for i in range(len(ks)):
    #     for line in lines:
    #         if abs(k_val(line) - ks[i]) <= k_d:
    #             ks_select[i].append(line)
    # ks_point = []
    # k_i_selected = []
    # for k_select in ks_select:
    #     k_x = []
    #     k_y = []
    #     for line in lines:
    #         k_x.append(line[0])
    #         k_x.append(line[2])
    #         k_y.append(line[1])
    #         k_y.append(line[3])
    #     k_x_avg = sum(k_x) / len(k_x)
    #     k_y_avg = sum(k_y) / len(k_y)
    #     ks_point.append([k_x_avg, k_y_avg])
    # 对于每条线，直接用这个斜率的直线淹没~
    draw = src.copy()
    if reach_cross:
        cv2.putText(draw, 'REACH_CROSS', (10, 100), cv2.FONT_HERSHEY_PLAIN, 1, (0, 255, 0), thickness=2)

    for t in theta_div2:
        k = np.tan(t / 360 * np.pi)
        logger.debug('k: %s' % k)
        for line in lines:
            d = abs(np.arctan(k_val(line)) * 360 / np.pi - t)
            if abs(np.arctan(k_val(line)) * 360 / np.pi - t) <= theta_d * 2:
                x, y = (line[0][0] + line[0][2]) / 2, (line[0][1] + line[0][3]) / 2
                pt1 = (0, int(y - k * x))
                pt2 = (draw.shape[1], int(k * (draw.shape[1] - x) + y))
                cv2.line(draw, pt1, pt2, (0, 0, 255), thickness=5)
    helper.Show.imshow(Configs.window_name + "绘图", draw, 1)
    helper.Controls.wait_exit(0)
    return src


def main():
    if not os.path.exists(Configs.file_path):
        raise FileNotFoundError
    for filename in os.listdir(Configs.file_path):
        if filename.split('.')[-1].lower() not in Configs.file_exts:
            continue
        logger.info("preparing %s..." % filename)
        src = cv2.imread(os.path.join(Configs.file_path, filename))
        logger.info("got image %s, %s" % (src.shape, src.dtype))
        logger.info('processing...')
        process(src)


if __name__ == '__main__':
    try:
        main()
    except Exception:
        traceback.print_exc()
        cv2.destroyAllWindows()
        helper.Controls.exit_kill()
